import os
import torch
import numpy as np
from copy import deepcopy
import torch.nn.functional as F

def save_ndarray(data_in, save_path, tensor_name='temp'):
    assert os.path.exists(save_path),'directory <{}> does not exist'.format(save_path)
    data = deepcopy(data_in)
    if isinstance(data_in, np.ndarray):
        pass
    elif isinstance(data_in, torch.Tensor):
        if data.device.type == 'cuda':
            data = data.detach().cpu()
        data = data.numpy()
    else:
        raise TypeError
    file = os.path.join(save_path, tensor_name+'.npy')
    np.save(file, data)
    print('data saved at {}, shape is <{}>\n'.format(file, data.shape))

def load_ndarray(load_path):
    print('load ndarray from {}'.format(load_path))
    assert os.path.exists(load_path),'directory <{}> does not exist'.format(load_path)
    return np.load(load_path)

def get_cosine_similarity(data_a, data_b):
    tensor_a = torch.tensor(data_a, dtype=torch.float32).view(1, -1)
    tensor_b = torch.tensor(data_b, dtype=torch.float32).view(1, -1)
    tensor_a = torch.clamp(tensor_a, -100, 100)
    tensor_b = torch.clamp(tensor_b, -100, 100)
    assert tensor_a.shape[1] == tensor_b.shape[1]
    similarity_score = F.cosine_similarity(tensor_a, tensor_b).item()
    return similarity_score

if __name__ == '__main__':
    # 创建一个示例数组
    array_int = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.int32)
    array_float = np.array([[20.9, 21.5, 31, 41], [51, 61, 71, 81]], dtype=np.float32)

    # 将数组保存为 .npy 文件
    np.save('python_int.npy', array_int)
    np.save('python_float.npy', array_float)



    # 从 .npy 文件中加载数组
    array_int2 = np.load('python_int.npy')
    array_float2 = np.load('/home/adt/models/wf_l-360/tensor/py/RPN_output_1_hei.npy')
    cplus_float2 = np.load('/home/adt/models/wf_l-360/tensor/cplus/RPN_output_1_hei.npy')


    # print(get_cosine_similarity(array_int, array_float))
    # print(get_cosine_similarity(array_int, array_int2))
    print(get_cosine_similarity(cplus_float2, array_float2))
    # print(get_cosine_similarity(array_int2, array_float2))